package cn.seecoder.ai.algorithm.ml.classification

import cn.seecoder.ai.algorithm.ml.BaseClassification
import cn.seecoder.ai.model.bo.TrainParamsBO
import ml.dmlc.xgboost4j.scala.spark.XGBoostClassifier
import org.apache.spark.ml.PipelineStage

class XGBoostClassification extends BaseClassification {


  override def buildMachineLearningStage(trainParams: TrainParamsBO): PipelineStage = {

    val xgb = new  XGBoostClassifier()
      .setFeaturesCol("indexedFeatures")
      .setLabelCol("label")
      .setMaxDepth(trainParams.maxDepth)
      .setMissing(0)




    if (trainParams.labelUniqueNum == 2) {
      xgb.setScalePosWeight(120)
    }
//     对于多分类，额外设置
    if(trainParams.labelUniqueNum>2){
      xgb.setNumClass(trainParams.labelUniqueNum)
    }

    xgb


  }


  }
